#include "PLNonLinearLeastSquaresSolver.h"

#include <iostream>

using std::cerr;
using std::cout;
using std::endl;
using std::string;
using std::vector;

double GetBiggestAbsValueFromVector( const vector<double> &myVector );
void HalfVectorValues( vector<double> *myVector );

/**
  *	Constructs a new PLNonLinearLeastSquaresSolver object with the basic options
  *
  * @param basicOptions basic Options of the PLNonLinearLeastSquaresSolver object
  */
PLNonLinearLeastSquaresSolver::PLNonLinearLeastSquaresSolver( const PLNonLinearLeastSquaresSolverOptions &basicOptions)
{
	mBasicOptions = basicOptions;
	mNumberOfVariables = 0;

	mLinearEquationSystem = 0;
}

/**
  *	Destroys the PLNonLinearLeastSquaresSolver object and releases the
  * memory needed for the matrix A and vectors x and b
  */
PLNonLinearLeastSquaresSolver::~PLNonLinearLeastSquaresSolver( void )
{
	if (mLinearEquationSystem) delete mLinearEquationSystem;
}

/**
  *	Returns the number of variables in the non linear equation system
  *
  * @return number of variables in the non linear equation system
  */
const PLuint &PLNonLinearLeastSquaresSolver::GetNumberOfVariables( void ) const
{
	return mNumberOfVariables;
}

/**
  *	Solves the non linear equation system for the current color channel
  * until one of the following condition holds:
  * 1) Maximum absolute gradient component below given threshold
  * 2) Line search fails
  * 3) Maximum number of iterations reached
  */
void PLNonLinearLeastSquaresSolver::Solve()
{
	double residual = Residual();

 	cout << endl << "Residual for the " << GetCurrentChannelName() << " channel before leveling: " << residual << endl << endl;

	PLuint iteration = 0;

	UINT64 colCount = mLinearEquationSystem->GetNumCols();
	UINT64 rowCount = mLinearEquationSystem->GetNumRows();

	vector<double> backup;
	vector<double> gradient;

	backup.resize( mLinearEquationSystem->GetNumRows(), 0.0 );
	gradient.resize( mLinearEquationSystem->GetNumRows(), 0.0 );

	for (; iteration < mBasicOptions.mMaxIterations && residual > mBasicOptions.mMinResidual; iteration++)
	{
		Reset();
		BuildLinearEquationSystem();

		mLinearEquationSystem->Finalize();

		cout << "Solving linear equation system..." << endl;

		mLinearEquationSystem->Solve();

		// Copy the values of x to gradient
		for (PLuint row = 0; row < rowCount; row++)
		{
			gradient[row] = mLinearEquationSystem->X( row );
		}

		cout << "Linear equation system solved. ";

		double maxAbsComponent = GetBiggestAbsValueFromVector( gradient );

		BackupVariables( &backup );

		PLuint lineSearchStep = 0;
		for (; lineSearchStep < mBasicOptions.mMaxLineSearchSteps && maxAbsComponent > mBasicOptions.mMinGradientAbsComponent; lineSearchStep++)
		{
			UpdateVariables( gradient );
			double r = Residual();

			if (r < residual)
			{
				residual = r;
				break;
			}
			else
			{
				RestoreVariables( backup );             
				HalfVectorValues( &gradient );
				maxAbsComponent = GetBiggestAbsValueFromVector( gradient );
			}
		}

		if (maxAbsComponent <= mBasicOptions.mMinGradientAbsComponent)
		{
			cerr << "Termination reason: Maximum absolute gradient component below given threshold" << endl;
			break;
		}

		if (lineSearchStep == mBasicOptions.mMaxLineSearchSteps)
		{
			cerr << "Termination reason: Line search failed" << endl;
			break;
		}

		for (; lineSearchStep < mBasicOptions.mMaxLineSearchSteps && lineSearchStep > 0; lineSearchStep++)
		{
			BackupVariables( &backup );
			HalfVectorValues( &gradient );
			UpdateVariables( gradient );

			double r = Residual();

			if (r < residual)
			{
				residual = r;
			}
			else
			{
				RestoreVariables( backup );
				break;
			}
		}
		cout << "Current residual is " << Residual() << endl;
	}

	cout << "Residual for the " << GetCurrentChannelName() << " channel at the end: " << residual << endl << endl;
	FinalizeSolver();
}

// Called line per line for the matrix J
/**
  *	Fills the matrix JTJ (=A) for a single row of J
  *
  * @param variables indexes of the variables in the current row of the matrix J
  * @param coefficients values of the variables in the row of J, which means that in the current row of J the variable at index variables[i] has the value coefficients[i]
  * @param weight weight of the variables
  * @param f function value corresponding to the current row of J
  */
void PLNonLinearLeastSquaresSolver::FillJTJ( const std::vector<PLuint> &variables,
	const std::vector<double> &coefficients,
	const double &weight, const double &f )
{
	mLinearEquationSystem->FillA( variables, coefficients, weight, f );
}

/**
  *	Resets the matrix A and vectors x and b to contain just 0-values
  */
void PLNonLinearLeastSquaresSolver::Reset()
{
	mLinearEquationSystem->Reset();
}

/**
  *	Initializes the matrix A and vectors x and b, which means that the needed
  * memory is allocated and the Reset() method is called
  */
void PLNonLinearLeastSquaresSolver::Initialize()
{
	if (mLinearEquationSystem) delete mLinearEquationSystem;
	mLinearEquationSystem = new PLLinearEquationSystem( mNumberOfVariables, mNumberOfVariables );
}

/**
  *	Returns the biggest absolute value of a vector
  *
  * @param myVector const reference to a vector with double values
  * @return biggest absolute value of the vector
  */
double GetBiggestAbsValueFromVector( const vector<double> &myVector )
{
	if (myVector.size() < 1)
		return 0.0;

	double biggest_abs_value = myVector[0];
	double current_value = 0.0;

	for (vector<double>::const_iterator it = myVector.begin()++; it != myVector.end(); it++)
	{
		current_value = *it > 0.0 ? *it : -*it;

		if (current_value > biggest_abs_value)
			biggest_abs_value = current_value;
	}

	return biggest_abs_value;
}

/**
  *	Halves all values of a vector
  *
  * @param myVector pointer to a vector with double values, that are halved
  */
void HalfVectorValues( vector<double> *myVector )
{
	if (myVector == 0) return;

	for (vector<double>::iterator it = myVector->begin(); it != myVector->end(); it++)
	{
		*it *= 0.5;
	}
}